# -*- coding: utf-8 -*-
import numpy as np
import scipy.io as sio
from sklearn.svm import SVC
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA


def classifier_train(train_x, train_y, classifier_type, filter_bank=False):
    """
    训练LDA/SVM，生成分类模型

    输入参数
    ----------
    train_x: array-like, shape (n_samples, n_features)
            L×2m   空间滤波后的数据samples
            L: 训练数据trial总数  m: processor 阶数
    train_y: array-like, shape (n_samples,)
            L 个trial对应的标签
    classifier_type: 训练模型 svm 或者lda

    返回值
    ----------
    svmModel: object

    """
    if filter_bank:
        # trial×filter×feature
        filter_num = train_x.shape[1]
        classifier_model = np.zeros(filter_num)
        for i in range(filter_num):
            if classifier_type == 'svm':
                classifier_model[i] = SVC(kernel='rbf', probability=True)
                classifier_model[i].fit(train_x[:, i, :], train_y)
            elif classifier_type == 'lda':
                classifier_model[i] = LDA()
                classifier_model[i].fit(train_x[:, i, :], train_y)
    else:
        if classifier_type == 'svm':
            classifier_model = SVC(kernel='rbf', probability=True)
            classifier_model.fit(train_x, train_y)
        elif classifier_type == 'lda':
            classifier_model = LDA()
            classifier_model.fit(train_x, train_y)
    return classifier_model


def classifier_predict(model, test_x, filter_bank=False):
    """
    根据SVM分类模型对测试数据进行分类

    输入参数
    ----------
    model: object
    test_x: array-like, shape (n_samples, n_features)
           L×2m   空间滤波后的数据
           L: 训练数据trial总数  m: processor 阶数
    返回值
    ----------
    predict_y: 1 维 L 个
              L个trial对应的预测标签

    """
    if filter_bank:
        # filter×feature
        filter_num = test_x.shape[0]
        predict_y = np.zeros(filter_num)
        for i in range(filter_num):
            predict_y[i] = model[i].predict(test_x[i, :])
    else:
        predict_y = model.predict(test_x)

    return predict_y


if __name__ == '__main__':
    dataForClassification = sio.loadmat(r'D:\Myfiles\WorkSpace\Codes\PythonProjects\Data\dataForClassification.mat')
    train_x = dataForClassification['train_x']  # shape(60,6)
    train_y = dataForClassification['train_y']  # shape(60,1)
    test_x = dataForClassification['test_x']  # shape(78,6)
    test_y = dataForClassification['test_y']  # shape(78,1)
    predict_y = dataForClassification['predict_y']  # shape(78,1)
    model = classifier_train(train_x, train_y.ravel(), 'svm')  # y.ravel()将2D(shapes, 1)改成1D(shapes, )的形式
    predict = classifier_predict(model, test_x)
    t = (predict == np.transpose(test_y))
    right = sum(map(sum, t))
    Accuracy = right / test_y.size
    print('正确率:', str(Accuracy))